{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "# 忽略烦人的红色提示\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "# 有 GPU 就用 GPU，没有就用 CPU\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "print('device', device)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from torchvision import transforms\n",
    "transform = transforms.Compose([transforms.Resize((64, 64)),\n",
    "                                    transforms.ToTensor(),\n",
    "                                    ]) #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])"
   ],
   "id": "4f42435170459e6f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "model = torch.load(r'D:\\Code\\2-ZSL\\Zero-Shot-Learning\\HSA-ZSL\\yms_zsl\\train\\output1\\models\\decae.pt',\n",
    "                   map_location='cpu', weights_only=False).to(device)"
   ],
   "id": "789fff5be780f06a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "model",
   "id": "949ebc6ad6f6155c",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from torchvision.models.feature_extraction import create_feature_extractor\n",
    "target_layers = {\n",
    "    'enc_block3.bn2': 'classifier_0',\n",
    "}\n",
    "model_trunc = create_feature_extractor(model, return_nodes=target_layers).to(device)\n",
    "model_trunc"
   ],
   "id": "3e42dba8b3c58e92",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "img_path = r'D:\\Code\\2-ZSL\\Zero-Shot-Learning\\data\\B\\seen\\val\\0-B-007_11.jpg'\n",
    "img_pil = Image.open(img_path)\n",
    "input_img = transform(img_pil) # 预处理\n",
    "input_img = input_img.unsqueeze(0).to(device)\n",
    "# 执行前向预测，得到指定中间层的输出\n",
    "pred_logits = model_trunc(input_img)"
   ],
   "id": "6eec77f48357c768",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "pred_logits['classifier_0'].squeeze().detach().cpu().numpy().shape",
   "id": "87bf040df7bc996",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "df = pd.read_csv(r'D:\\Code\\2-ZSL\\Zero-Shot-Learning\\MA-ZSL\\vae.csv', encoding=\"utf-8-sig\")",
   "id": "144fd1096d2f7ced",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "df.head()",
   "id": "2f7113dcb3c91a7b",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "df['标注类别名称']",
   "id": "1b2796c76bd4b2f0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "target_part = '0'\n",
    "part = 0\n",
    "sample = df[df['标注类别名称'].str.split('-', expand=True)[part] == target_part]\n",
    "sample"
   ],
   "id": "98b47bb485e419fe",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "encoding_array = []\n",
    "img_path_list = []\n",
    "model.eval()\n",
    "for img_path in tqdm(df['图像路径']):\n",
    "    img_path_list.append(img_path)\n",
    "    img_pil = Image.open(img_path)\n",
    "    input_img = transform(img_pil).unsqueeze(0).to(device) # 预处理\n",
    "    feature = model_trunc(input_img)['classifier_0'].view(-1).squeeze().detach().cpu().numpy() # 执行前向预测，得到 avgpool 层输出的语义特征\n",
    "    encoding_array.append(feature)\n",
    "encoding_array = np.array(encoding_array)"
   ],
   "id": "6c22aee846fd45f8",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "encoding_array.shape",
   "id": "f96513e111f5b463",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "np.save(r'D:\\Code\\2-ZSL\\Zero-Shot-Learning\\MA-ZSL\\特征提取.npy', encoding_array)",
   "id": "a464df6fef02e74c",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
